"""
使用训练得到的模型做预测
时间：2024/8/27 下午5:36
"""
import torch
import torch.nn as nn
from transformers import BertTokenizer, BertModel


class BertClassifier(nn.Module):
    '''
    定义了一个继承自 `torch.nn.Module` 的类 `BertClassifier`，用于文本分类任务。

    `__init__` 方法初始化BERT模型和分类器。
    `forward` 方法定义了模型的前向传播过程，使用BERT模型提取文本特征，并通过分类器进行分类。
    '''
    def __init__(self, bert_model, num_classes):
        """
        初始化方法接收预训练的BERT模型和分类类别数量作为参数。

        :param bert_model: `transformers` 中的预训练BERT模型，用于提取文本特征。
        :param num_classes: int，分类任务的目标类别数量。
        """
        super(BertClassifier, self).__init__()  # 调用父类的初始化方法
        self.bert = bert_model  # BERT模型用于提取文本特征
        self.classifier = nn.Linear(768, num_classes)   # 分类器用于将特征映射到类别

    def forward(self, input_ids, attention_mask):
        """
        定义了模型的前向传播过程。

        :param input_ids: 输入文本的token ID序列，形状为 (batch_size, sequence_length)。
        :param attention_mask: 输入文本的注意力掩码，形状为 (batch_size, sequence_length)，
        用于指示哪些位置是填充的。
        :return: 输出的类别概率，形状为 (batch_size, num_classes)。
        """
        # 使用BERT模型提取文本特征
        # 注意：BERT的输出包括序列的所有token的输出和pooled_output
        # 这里我们只关心pooled_output，不关心_，它通常用于分类任务
        pooled_output = self.bert(input_ids=input_ids,
                                  attention_mask=attention_mask).pooler_output
        # print(output)
        # pooled_output = output.pooler_output
        return self.classifier(pooled_output)


# 加载模型状态字典
model_path = "G:/HUGGINGFACE_HUB_CACHE/bert-base-uncased"
tokenizer = BertTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = BertModel.from_pretrained(model_path, trust_remote_code=True)

# 初始化分类模型
num_classes = 2
classifier = BertClassifier(model, num_classes)

# 加载模型状态字典
classifier.load_state_dict(torch.load('../model/model_state_dict.pth'))

# 将模型转移到相应的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
classifier.to(device)

# 设置模型为评估模式
classifier.eval()

# 新的文本数据
new_text = "I love this movie."

# 使用相同的 tokenizer 对文本进行编码
max_len = 64  # 使用与训练时相同的最大长度
encoding = tokenizer.encode_plus(
    new_text,
    add_special_tokens=True,
    max_length=max_len,
    padding='max_length',
    truncation=True,
    return_attention_mask=True,
    return_tensors='pt'
)

# 将编码后的数据转移到相应的设备
input_ids = encoding['input_ids'].to(device)
attention_mask = encoding['attention_mask'].to(device)

with torch.no_grad():
    # 使用模型进行前向传播
    outputs = classifier(input_ids, attention_mask)

# 获取预测结果
_, predicted_class = torch.max(outputs, dim=1)
predicted_label = predicted_class.item()

print(f"Predicted Label: {predicted_label}")
